"""
警告抑制器 - 用于抑制K-Means聚类中的常见警告
"""

import warnings
import contextlib
from functools import wraps

class KMeansWarningSupressor:
    """K-Means警告抑制器"""
    
    def __init__(self):
        self.suppressed_warnings = [
            "Number of distinct clusters",
            "Possibly due to duplicate points",
            "ConvergenceWarning",
            "n_clusters",
            "found smaller than"
        ]
    
    @contextlib.contextmanager
    def suppress_kmeans_warnings(self):
        """上下文管理器：抑制K-Means相关警告"""
        with warnings.catch_warnings():
            # 抑制所有sklearn相关警告
            warnings.filterwarnings("ignore", category=UserWarning)
            warnings.filterwarnings("ignore", category=FutureWarning)
            
            # 抑制特定的K-Means警告
            for warning_text in self.suppressed_warnings:
                warnings.filterwarnings("ignore", message=f".*{warning_text}.*")
            
            # 抑制sklearn的ConvergenceWarning
            try:
                from sklearn.exceptions import ConvergenceWarning
                warnings.filterwarnings("ignore", category=ConvergenceWarning)
            except ImportError:
                pass
            
            yield
    
    def suppress_warnings_decorator(self, func):
        """装饰器：为函数抑制K-Means警告"""
        @wraps(func)
        def wrapper(*args, **kwargs):
            with self.suppress_kmeans_warnings():
                return func(*args, **kwargs)
        return wrapper

# 全局实例
warning_suppressor = KMeansWarningSupressor()

# 便捷函数
def suppress_kmeans_warnings():
    """便捷函数：返回警告抑制上下文管理器"""
    return warning_suppressor.suppress_kmeans_warnings()

def kmeans_safe(func):
    """装饰器：为函数抑制K-Means警告"""
    return warning_suppressor.suppress_warnings_decorator(func)

# 使用示例
if __name__ == "__main__":
    print("警告抑制器测试")
    
    # 方法1：使用上下文管理器
    with suppress_kmeans_warnings():
        print("在此块中K-Means警告将被抑制")
    
    # 方法2：使用装饰器
    @kmeans_safe
    def test_function():
        print("此函数中的K-Means警告将被抑制")
    
    test_function()
    print("测试完成")
